# -*- coding: utf-8 -*-
"""
Elasticsearch客户端封装模块
提供与Elasticsearch交互的功能
"""

import time
from typing import Dict, List, Any, Optional, Generator
from elasticsearch import Elasticsearch, helpers
from elasticsearch.exceptions import ConnectionError, NotFoundError
from common.logger import get_logger
from common.config import config

logger = get_logger(__name__)


class ESClient:
    """Elasticsearch客户端封装类"""
    
    def __init__(self):
        """初始化ES客户端"""
        self.es_config = config.get('elasticsearch', {})
        self.client = self._create_client()
        self.index_pattern = self.es_config.get('index_pattern', 'logs-*')
        self.batch_size = self.es_config.get('batch_size', 1000)
        self.scroll_timeout = self.es_config.get('scroll_timeout', '5m')
    
    def _create_client(self) -> Elasticsearch:
        """
        创建Elasticsearch客户端
        
        Returns:
            Elasticsearch客户端实例
        """
        hosts = self.es_config.get('hosts', ['localhost:9200'])
        
        try:
            client = Elasticsearch(
                hosts=hosts,
                timeout=30,
                max_retries=3,
                retry_on_timeout=True
            )
            
            # 测试连接
            if client.ping():
                logger.info(f"成功连接到Elasticsearch: {hosts}")
                return client
            else:
                raise ConnectionError("无法连接到Elasticsearch")
                
        except Exception as e:
            logger.error(f"连接Elasticsearch失败: {e}")
            raise
    
    def search_logs(self, query: Dict[str, Any], start_time: Optional[str] = None, 
                   end_time: Optional[str] = None, size: Optional[int] = None) -> Generator[Dict[str, Any], None, None]:
        """
        搜索日志
        
        Args:
            query: 搜索查询
            start_time: 开始时间
            end_time: 结束时间
            size: 返回结果数量限制
            
        Yields:
            日志文档
        """
        # 构建时间范围查询
        if start_time or end_time:
            time_range = {}
            if start_time:
                time_range['gte'] = start_time
            if end_time:
                time_range['lte'] = end_time
            
            if 'bool' not in query:
                query = {'bool': {'must': [query]}}
            
            query['bool']['must'].append({
                'range': {
                    '@timestamp': time_range
                }
            })
        
        # 构建搜索体
        search_body = {
            'query': query,
            'sort': [{'@timestamp': {'order': 'asc'}}],
            'size': size or self.batch_size
        }
        
        try:
            # 使用scroll API进行分页查询
            response = self.client.search(
                index=self.index_pattern,
                body=search_body,
                scroll=self.scroll_timeout
            )
            
            scroll_id = response['_scroll_id']
            hits = response['hits']['hits']
            
            # 返回第一批结果
            for hit in hits:
                yield self._format_log_document(hit)
            
            # 继续滚动获取剩余结果
            while len(hits) > 0:
                response = self.client.scroll(
                    scroll_id=scroll_id,
                    scroll=self.scroll_timeout
                )
                
                scroll_id = response['_scroll_id']
                hits = response['hits']['hits']
                
                for hit in hits:
                    yield self._format_log_document(hit)
            
            # 清理scroll
            self.client.clear_scroll(scroll_id=scroll_id)
            
        except Exception as e:
            logger.error(f"搜索日志失败: {e}")
            raise
    
    def get_recent_logs(self, minutes: int = 5, size: int = 1000) -> List[Dict[str, Any]]:
        """
        获取最近的日志
        
        Args:
            minutes: 时间范围（分钟）
            size: 返回结果数量
            
        Returns:
            日志列表
        """
        end_time = time.time() * 1000  # 转换为毫秒
        start_time = end_time - (minutes * 60 * 1000)
        
        query = {
            'range': {
                '@timestamp': {
                    'gte': start_time,
                    'lte': end_time
                }
            }
        }
        
        logs = []
        for log in self.search_logs(query, size=size):
            logs.append(log)
        
        return logs
    
    def get_logs_by_app(self, app_name: str, minutes: int = 5, size: int = 1000) -> List[Dict[str, Any]]:
        """
        根据应用名称获取日志
        
        Args:
            app_name: 应用名称
            minutes: 时间范围（分钟）
            size: 返回结果数量
            
        Returns:
            日志列表
        """
        end_time = time.time() * 1000
        start_time = end_time - (minutes * 60 * 1000)
        
        query = {
            'bool': {
                'must': [
                    {'term': {'app_name.keyword': app_name}},
                    {
                        'range': {
                            '@timestamp': {
                                'gte': start_time,
                                'lte': end_time
                            }
                        }
                    }
                ]
            }
        }
        
        logs = []
        for log in self.search_logs(query, size=size):
            logs.append(log)
        
        return logs
    
    def get_logs_by_level(self, level: str, minutes: int = 5, size: int = 1000) -> List[Dict[str, Any]]:
        """
        根据日志级别获取日志
        
        Args:
            level: 日志级别（ERROR, WARN, INFO, DEBUG）
            minutes: 时间范围（分钟）
            size: 返回结果数量
            
        Returns:
            日志列表
        """
        end_time = time.time() * 1000
        start_time = end_time - (minutes * 60 * 1000)
        
        query = {
            'bool': {
                'must': [
                    {'term': {'level.keyword': level.upper()}},
                    {
                        'range': {
                            '@timestamp': {
                                'gte': start_time,
                                'lte': end_time
                            }
                        }
                    }
                ]
            }
        }
        
        logs = []
        for log in self.search_logs(query, size=size):
            logs.append(log)
        
        return logs
    
    def _format_log_document(self, hit: Dict[str, Any]) -> Dict[str, Any]:
        """
        格式化日志文档
        
        Args:
            hit: ES查询结果
            
        Returns:
            格式化的日志文档
        """
        source = hit.get('_source', {})
        return {
            'id': hit.get('_id'),
            'index': hit.get('_index'),
            'score': hit.get('_score'),
            'timestamp': source.get('@timestamp'),
            'message': source.get('message', ''),
            'level': source.get('level', 'INFO'),
            'app_name': source.get('app_name', ''),
            'host': source.get('host', ''),
            'source': source.get('source', ''),
            'raw_data': source
        }
    
    def get_indices(self) -> List[str]:
        """
        获取所有索引
        
        Returns:
            索引列表
        """
        try:
            response = self.client.cat.indices(format='json')
            return [index['index'] for index in response]
        except Exception as e:
            logger.error(f"获取索引列表失败: {e}")
            return []
    
    def get_index_stats(self, index_name: str) -> Dict[str, Any]:
        """
        获取索引统计信息
        
        Args:
            index_name: 索引名称
            
        Returns:
            索引统计信息
        """
        try:
            response = self.client.indices.stats(index=index_name)
            return response['indices'].get(index_name, {})
        except Exception as e:
            logger.error(f"获取索引统计信息失败: {e}")
            return {}
    
    def close(self):
        """关闭ES客户端连接"""
        if self.client:
            self.client.close() 